#include "CVehHook.h"

std::map<ULONG_PTR, VEHHOOK*> CVehHook::m_VehHookMap;

CVehHook::CVehHook()
{
    m_lpVeh = AddVectoredExceptionHandler(1, &CVehHook::VectoredExceptionHandler);
    if (!m_lpVeh)
        throw "Could not add Vectored Exception Handler";
}

CVehHook::~CVehHook()
{
    RemoveAll();

    if (m_lpVeh)
        RemoveVectoredExceptionHandler(m_lpVeh);
}

LONG NTAPI CVehHook::VectoredExceptionHandler(IN PEXCEPTION_POINTERS lpExceptionPointers)
{
#ifdef _WIN64
#define XIP Rip
#else
#define XIP Eip
#endif

    if (lpExceptionPointers->ExceptionRecord->ExceptionCode == EXCEPTION_BREAKPOINT)
    {
        for (auto it = m_VehHookMap.begin(); it != m_VehHookMap.end(); ++it)
        {
            if (lpExceptionPointers->ContextRecord->XIP != (ULONG_PTR)it->second->lpSrcFunc)
                continue;

            lpExceptionPointers->ContextRecord->XIP = (ULONG_PTR)it->second->lpNewFunc;
            return EXCEPTION_CONTINUE_EXECUTION;
        }
    }

    return EXCEPTION_CONTINUE_SEARCH;
}

BOOL CVehHook::AddVehHook(IN ULONG_PTR SrcFunc, IN ULONG_PTR NewFunc)
{
    PVEHHOOK lpVehHook = new VEHHOOK(SrcFunc, NewFunc);
    DWORD dwflOldProtect;
    BOOL bRet = FALSE;

    do
    {
        if (!lpVehHook) break;

        bRet = VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, PAGE_EXECUTE_READWRITE, &dwflOldProtect);
        if (!bRet) break;

        lpVehHook->OpCode = *((char*)lpVehHook->lpSrcFunc);
        *((char*)lpVehHook->lpSrcFunc) = '\xCC';

        bRet = VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, dwflOldProtect, &dwflOldProtect);
        if (!bRet) break;

        m_VehHookMap[(ULONG_PTR)SrcFunc] = lpVehHook;

        bRet = TRUE;

    } while (FALSE);

    if (!bRet && lpVehHook)
        delete lpVehHook;

    return bRet;
}

BOOL CVehHook::Remove(IN ULONG_PTR SrcFunc)
{
    DWORD dwflOldProtect;
    PVEHHOOK lpVehHook;
    BOOL bRet = FALSE;

    auto it = m_VehHookMap.find((ULONG_PTR)SrcFunc);
    if (it != m_VehHookMap.end())
    {
        lpVehHook = it->second;
        bRet = VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, PAGE_EXECUTE_READWRITE, &dwflOldProtect);
        if (bRet)
        {
            *((char*)lpVehHook->lpSrcFunc) = lpVehHook->OpCode;
            bRet = VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, dwflOldProtect, &dwflOldProtect);
            if (bRet)
            {
                delete lpVehHook;
                m_VehHookMap.erase(it);
            }
        }
    }

    return bRet;
}

void CVehHook::RemoveAll()
{
    DWORD dwflOldProtect;
    PVEHHOOK lpVehHook;

    for (auto it = m_VehHookMap.begin(); it != m_VehHookMap.end(); ++it)
    {
        lpVehHook = it->second;
        VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, PAGE_EXECUTE_READWRITE, &dwflOldProtect);
        *((char*)lpVehHook->lpSrcFunc) = lpVehHook->OpCode;
        VirtualProtect((LPVOID)lpVehHook->lpSrcFunc, 1, dwflOldProtect, &dwflOldProtect);
        delete it->second;
    }

    m_VehHookMap.clear();
    std::map<ULONG_PTR, VEHHOOK*>().swap(m_VehHookMap);
}
